import numpy as np
import torch
from torch import nn
from torch.nn import init
from collections import OrderedDict


class ProCBAMModule(nn.Module):
    class ChannelAttention(nn.Module):
        def __init__(self, channel, reduction=16):
            super().__init__()
            self.maxpool = nn.AdaptiveMaxPool2d(1)
            self.avgpool = nn.AdaptiveAvgPool2d(1)
            self.se = nn.Sequential(
                nn.Conv2d(channel, channel // reduction, 1, bias=False),
                nn.ReLU(),
                nn.Conv2d(channel // reduction, channel, 1, bias=False)
            )
            self.sigmoid = nn.Sigmoid()

        def forward(self, x):
            max_result = self.maxpool(x)
            avg_result = self.avgpool(x)
            max_out = self.se(max_result)
            avg_out = self.se(avg_result)
            output = self.sigmoid(max_out + avg_out)
            return output

    class SpatialAttention(nn.Module):
        def __init__(self, kernel_size=7):
            super().__init__()
            self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
            self.sigmoid = nn.Sigmoid()

        def forward(self, x):
            max_result, _ = torch.max(x, dim=1, keepdim=True)
            avg_result = torch.mean(x, dim=1, keepdim=True)
            result = torch.cat([max_result, avg_result], dim=1)
            output = self.conv(result)
            output = self.sigmoid(output)
            return output

    class SEAttention(nn.Module):

        def __init__(self, channel=512, reduction=16):
            super().__init__()
            self.avg_pool = nn.AdaptiveAvgPool2d(1)
            self.fc = nn.Sequential(
                nn.Linear(channel, channel // reduction, bias=False),
                nn.ReLU(inplace=True),
                nn.Linear(channel // reduction, channel, bias=False)
            )

        def init_weights(self):
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    init.kaiming_normal_(m.weight, mode='fan_out')
                    if m.bias is not None:
                        init.constant_(m.bias, 0)
                elif isinstance(m, nn.BatchNorm2d):
                    init.constant_(m.weight, 1)
                    init.constant_(m.bias, 0)
                elif isinstance(m, nn.Linear):
                    init.normal_(m.weight, std=0.001)
                    if m.bias is not None:
                        init.constant_(m.bias, 0)

        def forward(self, x):
            b, c, _, _ = x.size()
            y = self.avg_pool(x).view(b, c)
            y = self.fc(y).view(b, c, 1, 1)
            return y.expand_as(x)

    def __init__(self, channel, reduction=16, kernel_size=7, ):
        super().__init__()
        _empty = torch.empty(3, dtype=torch.float32, requires_grad=True)
        nn.init.constant_(_empty, 1 / 3)
        self.weights = nn.Parameter(_empty, requires_grad=True)
        self.ca = ProCBAMModule.ChannelAttention(channel=channel, reduction=reduction)
        self.sa = ProCBAMModule.SpatialAttention(kernel_size=kernel_size)
        self.se = ProCBAMModule.SEAttention(channel=channel, reduction=reduction)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        out_ca = x * self.ca(x)
        out_sa = x * self.sa(x)
        out_se = x * self.se(x)
        out = self.weights[0] * out_ca + self.weights[1] * out_sa + self.weights[2] * out_se
        return x + out


if __name__ == '__main__':
    device = "cuda"
    cs = ProCBAMModule(channel=512).to(device)
    x = torch.randn((2, 128, 14, 14))
    x = x.to(device)

    x = cs(x)
    pass
